package com.kara.woodAgent.agent.wood.node;

import com.alibaba.fastjson2.JSON;
import com.kara.woodAgent.agent.graph.Next;
import com.kara.woodAgent.agent.graph.Node;
import com.kara.woodAgent.agent.model.ModelProvider;
import com.kara.woodAgent.agent.wood.context.PlanContext;
import com.kara.woodAgent.agent.wood.context.WoodContext;
import com.kara.woodAgent.agent.wood.model.PlanResult;
import com.kara.woodAgent.agent.wood.model.SpecialAgentInfo;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.response.ChatResponse;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;


/**
 * @author tzy
 * @version 1.0
 * @project wood_agent
 * @description
 * @date 2025/6/10 10:00:07
 */

public class PlanningNode implements Node<WoodContext> {


	private List<SpecialAgentInfo> specialAgentInfoList;


	public PlanningNode(List<SpecialAgentInfo> specialAgentInfoList) {
		this.specialAgentInfoList = specialAgentInfoList;
	}

	@Autowired
	private ModelProvider modelProvider;

	/**
	 * #### 返回格式
	 * 请以json 数组 作为你的返回，数组中的每个对象都包含以下的字段：
	 * - id: 任务的唯一id (必须，从1开始，依次递增)
	 * - content 子任务的详细描述 (必须，包括子任务的具体内容)
	 * - agentName: 专项智能体的名称，对应 专项智能体的名称 字段 (必须)
	 * - status: 固定为 dispatch 字符串 (必须)
	 */


	private final String systemPrompt = """
			你是wood智能体系统的的子节点之一, 职责是：任务规划者。
			#### 任务
			你的任务是根据用户的任务描述 和 各个专项智能体的描述，规划出一个任务列表并 将任务分配给对应的专项智能体。
			#### 注意
			1. 请创建简洁、可执行且步骤清晰的任务列表。重点关注关键里程碑，而非详细的子步骤。优化任务列表以确保清晰高效。
			2. 一个任务只能分配给一个智能体。
			3. 任务分配遵循最小分配的原则，如果多个专项智能体间功能有重复，就尽量把任务整合分配给一个智能体
			#### 返回格式
			请以josn对象作为你的返回，包含一个字段：
				- plans: 任务列表 (必须),plans 是一个数组，每个元素是一个 任务对象
			&&&& 任务对象的格式
				- id: 任务的唯一id (必须，从1开始，依次递增)
				- content 子任务的详细描述 (必须，包括子任务的具体内容)
				- agentName: 专项智能体的名称 (必须)
				- status: 固定为 dispatch 字符串 (必须)
			#### 所有可用的专项智能体
			{{specialAgents}}
			""";



	@Override
	public Next execute(WoodContext context) {


		String specialAgents = specialAgentInfoList.stream().map(specialAgentInfo -> "专项智能体的名称: " +
				specialAgentInfo.agentName() +
				"  描述: " +
				specialAgentInfo.description()).collect(Collectors.joining("\n"));

		ChatLanguageModel chatModel = modelProvider.getChatModel();
		String userQuestion = context.getUserQuestion();
		List<ChatMessage> messages = List.of(
				SystemMessage.from(systemPrompt.replace("{{specialAgents}}", specialAgents)),
				UserMessage.from(userQuestion)
		);


		ChatRequest request = ChatRequest.builder().messages(messages).responseFormat(ResponseFormat.builder().type(ResponseFormatType.JSON).build())
				.build();
		ChatResponse chatResponse = chatModel.chat(request);

		AiMessage aiMessage = chatResponse.aiMessage();
		String text = aiMessage.text();

		PlanResult planResult = JSON.parseObject(text, PlanResult.class);
		context.setPlans(planResult.getPlans());
		PlanContext planContext = new PlanContext();
		ArrayList<ChatMessage> arrayList = new ArrayList<>(messages);
		arrayList.add(aiMessage);
		planContext.setPlanMessages(arrayList);
		context.setPlanContext(planContext);
		return Next.NextNode("dispatch").step(aiMessage.toString());
	}



}
